# In[]
import cv2
import numpy as np 
import os 
 
cap = cv2.VideoCapture(os.path.dirname(__file__) + '/vtest.avi')
# 角点检测参数
feature_params = dict(maxCorners = 500,
                        qualityLevel = 0.3,
                        minDistance = 7,
                        blockSize = 7)
# lk光流法参数
lk_params = dict(winSize = (15,15),
                maxLevel = 2,
                criteria = (cv2.TermCriteria_EPS | cv2.TermCriteria_COUNT,10,0.03))

# 计算第一帧特征点
ret,prev = cap.read()
preGray = cv2.cvtColor(prev,cv2.COLOR_BGR2GRAY)
p0 = cv2.goodFeaturesToTrack(preGray,mask=None,**feature_params)
count = 1
while True:
    ret,frame = cap.read()
    if not ret:
        break
    count+=1
    
    gray = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
    gray = cv2.equalizeHist(gray)
    # 计算光流
    p1,st,err = cv2.calcOpticalFlowPyrLK(preGray,gray,p0,None,**lk_params)
    #print(p0.shape,p1.shape,st.shape) #(56, 1, 2) (56, 1, 2) (56, 1)
    # 选取好的跟踪点
    goodPoints = p1[st==1]
    goodPrevPoints = p0[st==1]
    # print(goodPoints.shape) #(56, 2)
    # 在结果图像中叠加画出特征点和计算出的光流向量
    res = frame.copy()
    cv2.putText(res,str(count),(2,20),cv2.FONT_HERSHEY_PLAIN,1.5,color=[0,0,255])
    show = np.zeros_like(frame,dtype='float')
    redColor = [0,0,255]
    for i,(cur,prev) in enumerate(zip(goodPoints,goodPrevPoints)):
        x0,y0 = cur.ravel()
        x1,y1 = prev.ravel()
        dis = ((x0-x1)**2 + (y0-y1)**2)**0.5
        if dis < 2 or dis > 20:
            continue
        cv2.line(res,(x0,y0),(x1,y1),redColor)
        cv2.circle(res,(x0,y0),3,redColor)
        cv2.line(show,(x0,y0),(x1,y1),color=[255,255,255])

    # 更新上一帧
    preGray = gray.copy()
    p0 = goodPoints.reshape(-1,1,2)
    # print(p0.shape,'-'*100)

    # 显示计算结果图像
    cv2.imshow('video',res)
    cv2.imshow('show',show)

    if count in [150,161,545,763]:
        cv2.waitKey()
    c = cv2.waitKey(10)
    if c == -1: 
        pass
    elif chr(c) == ' ':
        cv2.waitKey()
    elif chr(c) == 'q':
        break
    